"""
算能 TPU Adapter 实现示例

此文件展示如何继承 BaseAdapter 实现自定义后端适配器
使用前需要：
1. 重命名为 sophon_tpu.py
2. 安装算能 TPU SDK（bmrt、SophonSDK 等）
3. 在 __init__.py 中导入此适配器
"""

from __future__ import annotations

import os
import subprocess
from pathlib import Path
from typing import Any, Dict, List

import numpy as np

# 注意：以下导入需要安装算能 TPU SDK
# import bmrt  # 算能 TPU Runtime API

from .base import AdapterBackend
from ..exceptions import ModelLoadError, InferenceError


class SophonTPUAdapter(AdapterBackend):
    """算能 TPU 后端适配器"""
    
    def __init__(self, device_id: int = 0):
        """
        初始化 Sophon TPU Adapter
        
        Args:
            device_id: TPU 设备 ID，默认为 0
        """
        self.device_id = device_id
        self.bmrt_handle = None
        self.bmodel_path = None
        self.input_names = None
        self.output_names = None
    
    def prepare(self, model_spec: Dict[str, Any]) -> Any:
        """
        加载 ONNX 模型并转换为 bmodel，然后加载到 TPU
        
        步骤：
        1. 校验模型合规性（SHA256）
        2. 将官方 ONNX 模型转换为 bmodel（使用 bmrt_convert）
        3. 加载 bmodel 到 TPU Runtime
        
        Args:
            model_spec: 模型规格字典
        
        Returns:
            bmrt.Runtime 句柄
        """
        try:
            # 导入 TPU SDK（延迟导入，避免在没有 TPU 的环境中失败）
            import bmrt
            
            # Step 1: 校验模型合规性
            self.validate_model_compliance(model_spec)
            
            source_onnx = model_spec.get("source_onnx")
            if not source_onnx:
                raise ModelLoadError("Missing 'source_onnx' in model_spec")
            
            # 检查文件是否存在
            if not Path(source_onnx).exists():
                raise ModelLoadError(f"Model file not found: {source_onnx}")
            
            # Step 2: 转换 ONNX → bmodel（如果尚未转换）
            self.bmodel_path = str(Path(source_onnx).with_suffix(".bmodel"))
            
            if not os.path.exists(self.bmodel_path):
                print(f"[INFO] 转换 ONNX → bmodel: {self.bmodel_path}")
                
                # 调用 bmrt_convert 工具进行转换
                # 注意：转换工具会保证模型拓扑、算子、IO 规格一致
                result = subprocess.run([
                    "bmrt_convert",
                    "--model", source_onnx,
                    "--output", self.bmodel_path,
                    "--target", "BM1684X",  # TPU 芯片型号（根据实际硬件调整）
                    "--opt", "2"  # 优化级别
                ], capture_output=True, text=True, check=True)
                
                print(f"[INFO] 转换成功: {result.stdout}")
            else:
                print(f"[INFO] 使用已存在的 bmodel: {self.bmodel_path}")
            
            # Step 3: 加载 bmodel 到 TPU Runtime
            self.bmrt_handle = bmrt.Runtime(self.device_id)
            self.bmrt_handle.load(self.bmodel_path)
            
            # Step 4: 获取模型输入输出信息
            self.input_names = list(model_spec["inputs"].keys())
            self.output_names = model_spec["outputs"]
            
            print(f"[INFO] 模型已加载到 TPU {self.device_id}")
            print(f"[INFO] 输入节点: {self.input_names}")
            print(f"[INFO] 输出节点: {self.output_names}")
            
            return self.bmrt_handle
            
        except subprocess.CalledProcessError as e:
            raise ModelLoadError(
                f"ONNX → bmodel 转换失败:\n"
                f"stdout: {e.stdout}\n"
                f"stderr: {e.stderr}"
            ) from e
        except Exception as e:
            if isinstance(e, ModelLoadError):
                raise
            raise ModelLoadError(f"Failed to prepare model on TPU: {e}") from e
    
    def infer(self, handle: Any, inputs: Dict[str, np.ndarray]) -> List[np.ndarray]:
        """
        使用 TPU 执行推理
        
        步骤：
        1. 将 numpy 输入转换为 TPU Tensor
        2. 在 TPU 上执行推理（调用 bmodel 中编译好的 TPU 算子）
        3. 将 TPU Tensor 输出转换回 numpy
        
        Args:
            handle: prepare() 返回的 bmrt.Runtime 句柄
            inputs: 输入张量字典 {input_name: np.ndarray}
        
        Returns:
            输出张量列表 List[np.ndarray]
        """
        try:
            import bmrt
            
            # Step 1: 准备输入数据（numpy → TPU Tensor）
            input_tensors = []
            for name in self.input_names:
                if name not in inputs:
                    raise InferenceError(f"Missing input: {name}")
                
                data = inputs[name]  # numpy array
                tensor = bmrt.Tensor(handle, data)  # 转换为 TPU Tensor
                input_tensors.append(tensor)
            
            # Step 2: 在 TPU 上执行推理（调用底层 TPU 算子）
            # 这里 handle.inference() 会执行 bmodel 中的 TPU 指令
            output_tensors = handle.inference(input_tensors)
            
            # Step 3: 获取输出结果（TPU Tensor → numpy）
            outputs = []
            for tensor in output_tensors:
                outputs.append(tensor.asnumpy())
            
            return outputs
            
        except Exception as e:
            raise InferenceError(f"TPU inference failed: {e}") from e
    
    def teardown(self, handle: Any) -> None:
        """
        释放 TPU 资源
        
        Args:
            handle: prepare() 返回的 bmrt.Runtime 句柄
        """
        if handle:
            try:
                handle.free()
            except Exception:
                pass  # 忽略释放失败
        
        self.bmrt_handle = None
        print("[INFO] TPU 资源已释放")
    
    def get_metadata(self) -> Dict[str, Any]:
        """
        返回 TPU 适配器和后端信息
        
        Returns:
            元信息字典
        """
        try:
            import bmrt
            backend_version = bmrt.__version__
        except Exception:
            backend_version = "unknown"
        
        return {
            "adapter_backend": "sophon_tpu",
            "adapter_version": "1.0.0",
            "backend": "SophonTPU",
            "backend_version": backend_version,
            "device": f"TPU:{self.device_id}",
            "providers": ["BMRuntimeExecutionProvider"],
            "conversion_tool": "bmrt_convert",
            "bmodel_path": self.bmodel_path,
            "optimization_flags": ["opt_level_2"]
        }

